package rs.pnv.util.xmlsec;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLConnection;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.Security;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.util.Collection;
import java.util.Iterator;

import javax.crypto.KeyGenerator;
import javax.crypto.SecretKey;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.FactoryConfigurationError;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerConfigurationException;
import javax.xml.transform.TransformerException;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.TransformerFactoryConfigurationError;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;

import org.apache.xml.security.encryption.EncryptedData;
import org.apache.xml.security.encryption.EncryptedKey;
import org.apache.xml.security.encryption.XMLCipher;
import org.apache.xml.security.encryption.XMLEncryptionException;
import org.apache.xml.security.keys.KeyInfo;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

public class Encrypt {
	
	
	
	static {
    	//staticka inicijalizacija
        Security.addProvider(new BouncyCastleProvider());
        org.apache.xml.security.Init.init();
    }
	
	/**
	 * Ulazi xml
	 */
	public static final String IN_FILE = "./data/mt102/mt102_signed.xml";
	/**
	 * Izlazini kriptovani xml
	 */
	public static final String OUT_FILE = "./data/mt102/mt102_enc_signed.xml";
	
	/**
	 * Cetifikat koji se cita
	 */
	public static final String KEY_STORE_FILE = "./cer/FirmaA.jks";
	
	
	/**
	 * Kreira DOM od XML dokumenta
	 */
	public static Document loadDocument(String file) {
		try {
			DocumentBuilderFactory dbf = DocumentBuilderFactory.newInstance();
			dbf.setNamespaceAware(true);
			DocumentBuilder db = dbf.newDocumentBuilder();
			Document document = db.parse(new File(file));

			return document;
		} catch (FactoryConfigurationError e) {
			e.printStackTrace();
			return null;
		} catch (ParserConfigurationException e) {
			e.printStackTrace();
			return null;
		} catch (SAXException e) {
			e.printStackTrace();
			return null;
		} catch (IOException e) {
			e.printStackTrace();
			return null;
		}
	}
	
	/**
	 * Ucitava sertifikat is KS fajla
	 * alias primer
	 */
	public static Certificate readCertificate(String pathToKeyStore,String password) {
		try {
			//kreiramo instancu KeyStore
			KeyStore ks = KeyStore.getInstance("JKS", "SUN");
			//ucitavamo podatke
			BufferedInputStream in = new BufferedInputStream(new FileInputStream(pathToKeyStore));
			ks.load(in, password.toCharArray());
			
			if(ks.isKeyEntry(password)) {
				Certificate cert = ks.getCertificate(password);
				return cert;
				
			}
			else
				return null;
			
		} catch (KeyStoreException e) {
			e.printStackTrace();
			return null;
		} catch (NoSuchProviderException e) {
			e.printStackTrace();
			return null;
		} catch (FileNotFoundException e) {
			e.printStackTrace();
			return null;
		} catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
			return null;
		} catch (CertificateException e) {
			e.printStackTrace();
			return null;
		} catch (IOException e) {
			e.printStackTrace();
			return null;
		} 
	}
	
	/**
	 * Snima DOM u XML fajl 
	 */
	public static void saveDocument(Document doc, String fileName) {
		try {
			File outFile = new File(fileName);
			FileOutputStream f = new FileOutputStream(outFile);

			TransformerFactory factory = TransformerFactory.newInstance();
			Transformer transformer = factory.newTransformer();
			
			DOMSource source = new DOMSource(doc);
			StreamResult result = new StreamResult(f);
			
			transformer.transform(source, result);

			f.close();

		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (TransformerConfigurationException e) {
			e.printStackTrace();
		} catch (IllegalArgumentException e) {
			e.printStackTrace();
		} catch (MalformedURLException e) {
			e.printStackTrace();
		} catch (TransformerFactoryConfigurationError e) {
			e.printStackTrace();
		} catch (TransformerException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	/**
	 * Generise tajni kljuc
	 */
	public static SecretKey generateDataEncryptionKey() {

        try {
			KeyGenerator keyGenerator = KeyGenerator.getInstance("AES"); //Triple DES
			return keyGenerator.generateKey();
		
        } catch (NoSuchAlgorithmException e) {
			e.printStackTrace();
			return null;
		}
    }
	
	
	/**
	 * Kriptuje sadrzaj prvog elementa odsek - odnosno onog elementa xml-a kojeg budemo naveli.
	 */
	public static Document encrypt(Document doc, SecretKey key, Certificate certificate,String elementName) {
		
		try {
			//cipher za kriptovanje tajnog kljuca,
			//Koristi se Javni RSA kljuc za kriptovanje
			XMLCipher keyCipher = XMLCipher.getInstance(XMLCipher.RSA_v1dot5);
		      //inicijalizacija za kriptovanje tajnog kljuca javnim RSA kljucem
		    keyCipher.init(XMLCipher.WRAP_MODE, certificate.getPublicKey());
		    EncryptedKey encryptedKey = keyCipher.encryptKey(doc, key);
			
		    //cipher za kriptovanje XML-a
		    XMLCipher xmlCipher = XMLCipher.getInstance(XMLCipher.AES_128);
		    //inicijalizacija za kriptovanje
		    xmlCipher.init(XMLCipher.ENCRYPT_MODE, key);
		    
		    //u EncryptedData elementa koji se kriptuje kao KeyInfo stavljamo kriptovan tajni kljuc
		    EncryptedData encryptedData = xmlCipher.getEncryptedData();
	        //kreira se KeyInfo
		    KeyInfo keyInfo = new KeyInfo(doc);
		    keyInfo.addKeyName("Kriptovani tajni kljuc");
	        //postavljamo kriptovani kljuc
		    keyInfo.add(encryptedKey);
		    //postavljamo KeyInfo za element koji se kriptuje
	        encryptedData.setKeyInfo(keyInfo);
			
			//trazi se element ciji sadrzaj se kriptuje
			NodeList odseci = doc.getElementsByTagName(elementName);
			Element odsek = (Element) odseci.item(0);
			
			xmlCipher.doFinal(doc, odsek, true); //kriptuje sa sadrzaj
			
			return doc;
			
		} catch (XMLEncryptionException e) {
			e.printStackTrace();
			return null;
		} catch (Exception e) {
			e.printStackTrace();
			return null;
		}
	}
	
	
//	@SuppressWarnings("rawtypes")
//	public static Certificate readCertificate(String path){
//		try {
//			FileInputStream fis = new FileInputStream(path);
//			CertificateFactory cf = CertificateFactory.getInstance("X.509");
//			Collection c = cf.generateCertificates(fis);
//			Iterator i = c.iterator();
//			while (i.hasNext()) {
//			    Certificate cert = (Certificate)i.next();
//			    return cert;
//			}
//		} catch (FileNotFoundException e) {
//			e.printStackTrace();
//		} catch (CertificateException e) {
//			e.printStackTrace();
//		}
//		return null;
//	}
	
	
	@SuppressWarnings("rawtypes")
	public static Certificate readRemoteCertificate(String cerPath){
		try {
			
			URL url = new URL(cerPath);
			URLConnection connection = url.openConnection();
			CertificateFactory cf = CertificateFactory.getInstance("X.509");
			Collection c = cf.generateCertificates(connection.getInputStream());
			Iterator i = c.iterator();
			while (i.hasNext()) {
			    Certificate cert = (Certificate)i.next();
			    return cert;
			}
		} catch (FileNotFoundException e) {
			e.printStackTrace();
		} catch (CertificateException e) {
			e.printStackTrace();
		} catch (MalformedURLException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		return null;
	}
	
	
	public Certificate readCerFromFile(String path) {

		@SuppressWarnings("rawtypes")
		Collection collection;
		CertificateFactory certificateFactory;

		collection = null;
		certificateFactory = null;

		try {

			InputStream inputStream = getClass().getResourceAsStream(path);
			certificateFactory = CertificateFactory.getInstance("X.509");
			collection = certificateFactory.generateCertificates(inputStream);

			@SuppressWarnings("rawtypes")
			Iterator i = collection.iterator();
			while (i.hasNext()) {
				Certificate cert = (Certificate) i.next();
				return cert;
			}

		} catch (CertificateException e) {
			System.out.println(e.getMessage());
		} finally {

		}

		return null;
	}
	
	/** 
	 * @param docLocation - "./data/send/send.xml"
	 * @param cerLocation - "./cer/BankaA.cer"
	 * @param saveLocation - "./data/send/send_enc.xml"
	 */
	public Document docEncryption(Document doc,String cerLocation) {
		System.out.println("#####Kriptovanje#####");
		System.out.println("Generisanje tajnog kljuca.");
		SecretKey secretKey = generateDataEncryptionKey();
		System.out.println("Izvlacenje setifikata");
		//Certificate certificate = readRemoteCertificate(cerLocation);
		Certificate certificate  = readCerFromFile(cerLocation);
		System.out.println("Generisanje kriptovanog dokumenta.");		
		if(certificate == null) System.out.println("JEPSE CERTIFIKATIU!");
		System.out.println("Node to encrypt : " + doc.getDocumentElement().getNodeName());
		Document retDoc = encrypt(doc, secretKey, certificate,doc.getDocumentElement().getNodeName());
		System.out.println("#####Kriptovanje#####");
		return retDoc;
	}
	
//	public static void main(String[] args) {
//		System.out.println("Ucitavnje dokumenta.");
//		Document doc = loadDocument("./data/send/send.xml");
//		System.out.println("Generisanje tajnog kljuca.");
//		SecretKey secretKey = generateDataEncryptionKey();
//		System.out.println("Izvlacenje setifikata");
//		//Certificate certificate = readCertificate(EncyptUtil.KEY_STORE_FILE, "FirmaA");
//		Certificate certificate = readCertificate("./cer/BankaA.cer");
//		System.out.println("Generisanje kriptovanog dokumenta.");		
//		doc = encrypt(doc, secretKey, certificate,doc.getDocumentElement().getNodeName());
//		System.out.println("Snimanje kriptovanog kokumenta.");
//		saveDocument(doc, "./data/send/send_enc.xml");
//	
//	}

}
